import cv2
import time
import numpy as np

model_dir = 'model'
classesFile = model_dir+'/coco.names'
modelConfiguration = model_dir+'/yolov3.cfg'
modelWeights = model_dir+'/yolov3.weights'
video = 'C:/Users/tellw/Videos/BiliBiliDownload/7879930/12933648/000.mp4'
confThreshold=0.4
nmsThreshold=0.3
inpWidth=416
inpHeight=416
with open(classesFile, 'r') as f:
    classes = f.read().rstrip('\n').split('\n')
    print(classes)

def postprocess(frame, outs):
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]
    classIds = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            if confidence > confThreshold:
                center_x = int(detection[0]*frameWidth)
                center_y = int(detection[1]*frameHeight)
                width = int(detection[2]*frameWidth)
                height = int(detection[3]*frameHeight)
                left = int(center_x-width/2)
                top = int(center_y-height/2)
                classIds.append(classId)
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])
    indices = cv2.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
    for i in indices:
        i = i[0]
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        drawPred(classIds[i], confidences[i], left, top, left+width, top+height, frame)

def drawPred(classId, conf, left, top, right, bottom, frame):
    cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255))
    label = '%.2f'%conf
    if classes:
        assert classId < len(classes)
        label = '%s:%s'%(classes[classId], label)
    labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    top = max(top, labelSize[1])
    cv2.putText(frame, label, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))

def getOutputsNames(net):
    laysNames = net.getLayerNames()
    return [laysNames[i[0]-1] for i in net.getUnconnectedOutLayers()]

if __name__ == '__main__':
    choice = input('choose to detect originally by 1 or detect customly by 2')
    if choice == '1':
        net = cv2.dnn.readNetFromDarknet(modelConfiguration, modelWeights)
        net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
        net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
        cap = cv2.VideoCapture(video)
        while True:
            success, frame = cap.read()
            if success:
                startTime = time.time()
                blob = cv2.dnn.blobFromImage(frame, 1/255, (inpWidth, inpHeight), (0, 0, 0), 1, crop=False)
                net.setInput(blob)
                outs = net.forward(getOutputsNames(net))
                postprocess(frame, outs)
                label = 'inference time for frame: %.2f s'%(time.time()-startTime)
                cv2.putText(frame, label, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
                cv2.imshow('detect', frame)
                k = cv2.waitKey(40)
                if k == 27 or k == ord('q'):
                    break
            else:
                break
        cap.release()
        cv2.destroyAllWindows()
    else:
        pass
